"""
JSON 数据存储层
实现知识库数据的增删改查操作
"""
import json
import time
from pathlib import Path
from typing import List, Dict, Any, Optional
from threading import Lock


class KnowledgeBaseStorage:
    """知识库存储管理类"""
    
    def __init__(self, storage_path: Path, embedding_model: str):
        self.storage_path = Path(storage_path)
        self.embedding_model = embedding_model
        self._lock = Lock()  # 文件锁，确保并发安全
        self._ensure_storage_exists()
    
    def _ensure_storage_exists(self):
        """确保存储目录和文件存在"""
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        if not self.storage_path.exists():
            self._initialize_empty_kb()
    
    def _initialize_empty_kb(self):
        """初始化空知识库"""
        data = {
            "global_metadata": {
                "total_documents": 0,
                "total_chunks": 0,
                "created_at": time.time(),
                "updated_at": time.time(),
                "model": self.embedding_model
            },
            "documents": [],
            "chunks": []
        }
        self._write_data(data)
    
    def _read_data(self) -> Dict[str, Any]:
        """读取知识库数据"""
        with self._lock:
            try:
                with open(self.storage_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except (FileNotFoundError, json.JSONDecodeError):
                self._initialize_empty_kb()
                with open(self.storage_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
    
    def _write_data(self, data: Dict[str, Any]):
        """写入知识库数据"""
        with self._lock:
            with open(self.storage_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
    
    def add_document(
        self,
        filename: str,
        file_size: int,
        chunks_text: List[str],
        embeddings: List[List[float]],
        chunk_size: int
    ) -> int:
        """
        添加文档到知识库
        
        Args:
            filename: 文件名
            file_size: 文件大小
            chunks_text: 文本块列表
            embeddings: 向量列表
            chunk_size: 分块大小
            
        Returns:
            添加的文本块数量
        """
        data = self._read_data()
        
        # 如果文档已存在，先删除旧数据
        existing_doc = self._find_document(data, filename)
        if existing_doc:
            self._remove_document_chunks(data, existing_doc)
            data["documents"] = [d for d in data["documents"] if d["filename"] != filename]
        
        # 计算新的起始索引
        next_index = max([chunk["index"] for chunk in data["chunks"]], default=-1) + 1
        
        # 创建新的文本块
        new_chunks = []
        chunk_indices = []
        for i, (text, embedding) in enumerate(zip(chunks_text, embeddings)):
            chunk_index = next_index + i
            new_chunks.append({
                "index": chunk_index,
                "document": filename,
                "text": text,
                "embedding": embedding
            })
            chunk_indices.append(chunk_index)
        
        # 添加文档元数据
        total_characters = sum(len(text) for text in chunks_text)
        doc_metadata = {
            "filename": filename,
            "file_size": file_size,
            "upload_time": time.time(),
            "chunk_size": chunk_size,
            "chunk_count": len(chunks_text),
            "total_characters": total_characters,
            "chunk_indices": chunk_indices
        }
        
        # 更新数据
        data["chunks"].extend(new_chunks)
        data["documents"].append(doc_metadata)
        
        # 更新全局元数据
        data["global_metadata"]["total_documents"] = len(data["documents"])
        data["global_metadata"]["total_chunks"] = len(data["chunks"])
        data["global_metadata"]["updated_at"] = time.time()
        
        self._write_data(data)
        return len(new_chunks)
    
    def get_all_chunks(self) -> List[Dict[str, Any]]:
        """获取所有文本块"""
        data = self._read_data()
        return data["chunks"]
    
    def get_document_list(self) -> List[Dict[str, Any]]:
        """获取文档列表"""
        data = self._read_data()
        return data["documents"]
    
    def get_status(self) -> Dict[str, Any]:
        """获取知识库状态"""
        data = self._read_data()
        metadata = data["global_metadata"]
        
        # 计算总字符数
        total_characters = sum(doc["total_characters"] for doc in data["documents"])
        
        return {
            "total_documents": metadata["total_documents"],
            "total_chunks": metadata["total_chunks"],
            "total_characters": total_characters,
            "last_updated": metadata.get("updated_at"),
            "embedding_model": metadata["model"],
            "is_initialized": metadata["total_chunks"] > 0
        }
    
    def delete_document(self, filename: str) -> Optional[Dict[str, Any]]:
        """
        删除指定文档
        
        Args:
            filename: 要删除的文件名
            
        Returns:
            删除的文档信息，如果文档不存在返回 None
        """
        data = self._read_data()
        
        # 查找文档
        doc = self._find_document(data, filename)
        if not doc:
            return None
        
        # 删除文本块
        deleted_chunks = self._remove_document_chunks(data, doc)
        
        # 删除文档记录
        data["documents"] = [d for d in data["documents"] if d["filename"] != filename]
        
        # 更新全局元数据
        data["global_metadata"]["total_documents"] = len(data["documents"])
        data["global_metadata"]["total_chunks"] = len(data["chunks"])
        data["global_metadata"]["updated_at"] = time.time()
        
        self._write_data(data)
        
        return {
            "filename": filename,
            "deleted_chunks": deleted_chunks
        }
    
    def _find_document(self, data: Dict[str, Any], filename: str) -> Optional[Dict[str, Any]]:
        """查找文档"""
        for doc in data["documents"]:
            if doc["filename"] == filename:
                return doc
        return None
    
    def _remove_document_chunks(self, data: Dict[str, Any], doc: Dict[str, Any]) -> int:
        """删除文档的所有文本块"""
        chunk_indices_set = set(doc["chunk_indices"])
        original_count = len(data["chunks"])
        data["chunks"] = [chunk for chunk in data["chunks"] if chunk["index"] not in chunk_indices_set]
        return original_count - len(data["chunks"])

